import tensorflow as tf

X = tf.Variable([[[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [7, 8, 1]]], dtype=tf.float32);

mean_0 = tf.reduce_mean(X, axis=0)
mean_1 = tf.reduce_mean(X, axis=1)
mean_0_1 = tf.reduce_mean(X, axis=[1, 0])
mean_1_0 = tf.reduce_mean(X, axis=[0, 1])

d_mean_0 = tf.gradients(mean_0, X)
d_mean_1 = tf.gradients(mean_1, X)
d_mean_0_1 = tf.gradients(mean_0_1, X)
d_mean_1_0 = tf.gradients(mean_1_0, X)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    
    print('------mean_0------')
    print(sess.run(mean_0))
    print('------mean_1------')
    print(sess.run(mean_1))
    print('------mean_0_1------')
    print(sess.run(mean_0_1))
    print('------mean_1_0------')
    print(sess.run(mean_1_0))
    
    print('------d_mean_0------')
    print(sess.run(d_mean_0))
    print('------d_mean_1------')
    print(sess.run(d_mean_1))
    print('------d_mean_0_1------')
    print(sess.run(d_mean_0_1))
    print('------d_mean_1_0------')
    print(sess.run(d_mean_1_0))